import tensorflow as tf


# 2d卷积层
def conv2d(input_tensor, name_scope, output_channel, kh, kw, dh, dw, padding='SAME', activation=tf.nn.relu,
           training=True, batch_normalize=False, momentum=0.997, epsilon=1e-3):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param output_channel: 输出张量的通道数
    :param kh: 卷积核的宽度
    :param kw: 卷积核的高度
    :param dh: 卷积核的水平移动步长
    :param dw: 卷积核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :param activation: 激活函数的类型
    :param training: 是否可以被训练
    :param batch_normalize: 是否具有批处理归一化结构
    :param momentum: 仅当batch_normalize=True时生效，批处理归一化的参数
    :param epsilon: 仅当batch_normalize=True时生效，批处理归一化的参数
    :return: 输出的张量
    """
    print(input_tensor)
    input_channel = input_tensor.get_shape()[-1].value  # 获取输入张量的通道数，即图像的通道数或者feature map数量
    with tf.name_scope(name_scope) as scope:
        # 定义权重
        w = tf.get_variable(scope + 'w', shape=[kh, kw, input_channel, output_channel],
                            dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer_conv2d())
        # 卷积操作
        conv = tf.nn.conv2d(input_tensor, w, (1, dh, dw, 1), padding=padding)
        # 定义偏置
        b = tf.get_variable(scope + 'b', shape=[output_channel], dtype=tf.float32,
                            initializer=tf.constant_initializer(0.01))
        # 计算输出的张量
        z = tf.nn.bias_add(conv, b)
        # 批处理归一化
        if batch_normalize:
            z = tf.layers.batch_normalization(z, momentum=momentum, epsilon=epsilon, trainable=training)
        if activation:
            output_tensor = activation(z, name=scope)   # 激活
        else:
            output_tensor = z
        return output_tensor


# 2d反卷积层（转置卷积）
def deconv2d(input_tensor, name_scope, output_channel, kh, kw, dh, dw, padding='SAME', activation=tf.nn.relu):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param output_channel: 输出张量的通道数
    :param kh: 卷积核的宽度
    :param kw: 卷积核的高度
    :param dh: 卷积核的水平移动步长
    :param dw: 卷积核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :param activation: 激活函数的类型
    :return: 输出的张量
    """
    print(input_tensor)
    input_channel = input_tensor.get_shape()[-1].value  # 获取输入张量的通道数，即图像的通道数或者feature map数量
    with tf.name_scope(name_scope) as scope:
        # 定义权重
        w = tf.get_variable(scope + 'w', shape=[kh, kw, input_channel, output_channel],
                            dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer_conv2d())
        # 卷积操作
        conv = tf.nn.conv2d_transpose(input_tensor, w, (1, dh, dw, 1), padding=padding)
        if activation:
            output_tensor = activation(conv, name=scope)    # 激活
        else:
            output_tensor = conv
        return output_tensor


# 2d深度卷积
def depthwise_conv2d(input_tensor, name_scope, kh, kw, dh, dw, activation=tf.nn.relu, momentum=0.9997, epsilon=1e-3):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param kh: 卷积核的宽度
    :param kw: 卷积核的高度
    :param dh: 卷积核的水平移动步长
    :param dw: 卷积核的垂直移动步长
    :param activation: 激活函数的类型
    :param momentum: 仅当batch_normalize=True时生效，批处理归一化的参数
    :param epsilon: 仅当batch_normalize=True时生效，批处理归一化的参数
    :return: 输出的张量
    """
    print(input_tensor)
    input_channel = input_tensor.get_shape()[-1].value  # 获取输入张量的通道数，即图像的通道数或者feature map数量
    with tf.name_scope(name_scope) as scope:
        # 定义权重
        w = tf.get_variable(scope + 'w', shape=[kh, kw, input_channel, 1],
                            dtype=tf.float32, initializer=tf.zeros_initializer(), trainable=True)
        dw_conv = tf.nn.depthwise_conv2d(input_tensor, w, strides=[1, dh, dw, 1], padding='SAME', rate=[1, 1])
        # 批处理归一化
        dw_bn = tf.layers.batch_normalization(dw_conv, momentum=momentum, epsilon=epsilon, training=True)
        if activation:
            output_tensor = activation(dw_bn, name=scope)   # 激活
        else:
            output_tensor = dw_bn
        return output_tensor


# 3d卷积层
def conv3d(input_tensor, name_scope, output_channel, kd, kh, kw, dd, dh, dw, padding='SAME', activation=tf.nn.relu,
           training=True, batch_normalize=False, momentum=0.997, epsilon=1e-3):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param output_channel: 输出张量的通道数
    :param kd: 卷积核的深度
    :param kh: 卷积核的宽度
    :param kw: 卷积核的高度
    :param dd: 卷积核的纵向移动步长
    :param dh: 卷积核的水平移动步长
    :param dw: 卷积核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :param activation: 激活函数的类型
    :param training: 是否可以被训练
    :param batch_normalize: 是否具有批处理归一化结构
    :param momentum: 仅当batch_normalize=True时生效，批处理归一化的参数
    :param epsilon: 仅当batch_normalize=True时生效，批处理归一化的参数
    :return: 输出的张量
    """
    print(input_tensor)
    input_channel = input_tensor.get_shape()[-1].value  # 获取输入张量的通道数，即图像的通道数或者feature map数量
    with tf.name_scope(name_scope) as scope:
        # 定义权重
        w = tf.get_variable(scope + 'w', shape=[kd, kh, kw, input_channel, output_channel], dtype=tf.float32)
        # 卷积操作
        conv = tf.nn.conv3d(input_tensor, w, (1, dd, dh, dw, 1), padding=padding)
        # 定义偏置
        b = tf.get_variable(scope + 'b', shape=[output_channel], dtype=tf.float32,
                            initializer=tf.constant_initializer(0.01))
        # 计算输出的张量
        z = tf.nn.bias_add(conv, b)
        # 批处理归一化
        if batch_normalize:
            z = tf.layers.batch_normalization(z, momentum=momentum, epsilon=epsilon, trainable=training)
        if activation:
            output_tensor = activation(z, name=scope)
        else:
            output_tensor = z
        return output_tensor


# 2d最大池化层
def maxpool2d(input_tensor, name_scope, kh, kw, dh, dw, padding='SAME'):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param kh: 核的宽度
    :param kw: 核的高度
    :param dh: 核的水平移动步长
    :param dw: 核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :return: 输出的张量
    """
    print(input_tensor)
    return tf.nn.max_pool(input_tensor, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding=padding, name=name_scope)


# 3d最大池化层
def maxpool3d(input_tensor, name_scope, kd, kh, kw, dd, dh, dw, padding='SAME'):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param kd: 核的深度
    :param kh: 核的宽度
    :param kw: 核的高度
    :param dd: 核的纵向移动步长
    :param dh: 核的水平移动步长
    :param dw: 核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :return: 输出的张量
    """
    print(input_tensor)
    return tf.nn.max_pool3d(input_tensor,
                            ksize=[1, kd, kh, kw, 1], strides=[1, dd, dh, dw, 1], padding=padding, name=name_scope)


# 2d全局平均池化层
def avgpool2d(input_tensor, name_scope, kh, kw, dh, dw, padding='SAME'):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param kh: 核的宽度
    :param kw: 核的高度
    :param dh: 核的水平移动步长
    :param dw: 核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :return: 输出的张量
    """
    print(input_tensor)
    return tf.nn.avg_pool(input_tensor, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding=padding, name=name_scope)


# 3d全局平均池化层
def avgpool3d(input_tensor, name_scope, kd, kh, kw, dd, dh, dw, padding='SAME'):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param kd: 核的深度
    :param kh: 核的宽度
    :param kw: 核的高度
    :param dd: 核的纵向移动步长
    :param dh: 核的水平移动步长
    :param dw: 核的垂直移动步长
    :param padding: 可选SAME和VALID。例：张量宽10，卷积核宽3，水平步长3，则卷积核在滑动至张量边界时出现1个多余的数，不能直接进行卷积运算。
    取SAME时，将在张量末尾添加2个0补齐，并进行卷积运算；取VALID时，不添加0，直接省略最后的余数。
    :return: 输出的张量
    """
    print(input_tensor)
    return tf.nn.avg_pool3d(input_tensor,
                            ksize=[1, kd, kh, kw, 1], strides=[1, dd, dh, dw, 1], padding=padding, name=name_scope)


# 全连接层
def fc(input_tensor, name_scope, output_channel, activation=tf.nn.relu, regularizer=None):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :param output_channel: 输出张量的通道数
    :param activation: 激活函数的类型
    :param regularizer: 正则化项
    :return: 输出的张量
    """
    print(input_tensor)
    input_channel = input_tensor.get_shape()[-1].value  # 获取输入张量的通道数，即图像的通道数或者feature map数量
    with tf.name_scope(name_scope) as scope:
        w = tf.Variable(tf.truncated_normal([input_channel, output_channel],
                                            stddev=0.1, dtype=tf.float32), name=scope + 'w')
        # 进行正则化
        if regularizer is not None:
            tf.add_to_collection('losses', regularizer(w))
        b = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[output_channel]), name=scope + 'b')
        z = tf.matmul(input_tensor, w) + b
        if activation:
            output_tensor = activation(z, name=scope)   # 激活
        else:
            output_tensor = z
    return output_tensor
